"""
将整染色体 Hi-C 矩阵（.npy，shape=N×N）转换为 .cool 文件。

注意：
- 需要指定染色体名与分辨率；可选提供染色体长度（bp），用于设置最后一个 bin 的 end，未提供则按 (i+1)*resolution。
- 默认仅写上三角 (i<=j)，并设置 symmetric_upper=True。
"""
import os
import argparse
import numpy as np
import pandas as pd


def parse_args():
    parser = argparse.ArgumentParser(description="Convert N×N .npy Hi-C matrix to .cool")
    parser.add_argument("--npy", type=str, required=True, help="输入的 .npy 路径（float32/float64 矩阵，N×N）")
    parser.add_argument("--output", type=str, required=True, help="输出 .cool 路径")
    parser.add_argument("--chromosome", type=str, required=True, help="染色体名（如 chr1）")
    parser.add_argument("--resolution", type=int, required=True, help="分辨率（bp）")
    parser.add_argument("--chrom_size", type=int, default=0, help="可选：染色体长度（bp）")
    parser.add_argument("--min_count_threshold", type=float, default=0.0, help="像素最小阈值， <=阈值 的像素将被忽略")
    parser.add_argument("--select_idx", type=int, default=0, help="当 .npy 为 3D (B,H,W) 或 4D (B,C,H,W) 时，选择第几个样本/通道")
    parser.add_argument("--reduce", type=str, default="none", choices=["none", "mean", "max"], help="当 .npy 为 3D/4D 时，是否在批维上做聚合")
    args = parser.parse_args()
    return args


def build_bins_df(chrom: str, num_bins: int, resolution: int, chrom_size: int = 0) -> pd.DataFrame:
    starts = np.arange(0, num_bins * resolution, resolution, dtype=np.int64)
    if chrom_size and chrom_size > 0:
        ends = np.minimum(starts + resolution, chrom_size).astype(np.int64)
    else:
        ends = (starts + resolution).astype(np.int64)
    return pd.DataFrame({
        "chrom": [chrom] * num_bins,
        "start": starts,
        "end": ends,
    })


def build_pixels_df(matrix: np.ndarray, threshold: float = 0.0) -> pd.DataFrame:
    assert matrix.ndim == 2 and matrix.shape[0] == matrix.shape[1]
    n = matrix.shape[0]
    i, j = np.triu_indices(n)
    vals = matrix[i, j].astype(np.float64)
    if threshold > 0:
        mask = vals > threshold
        i, j, vals = i[mask], j[mask], vals[mask]
    return pd.DataFrame({
        "bin1_id": i.astype(np.int64),
        "bin2_id": j.astype(np.int64),
        "count": vals,
    })


def main():
    args = parse_args()
    arr = np.load(args.npy)
    # 兼容 2D/3D/4D：
    # - 2D: (H,W)
    # - 3D: (B,H,W) -> 选 idx 或聚合
    # - 4D: (B,C,H,W) -> 优先在 C 维压到单通道，再选/聚合 B 维
    if arr.ndim == 2:
        mat = arr
    elif arr.ndim == 3:
        if args.reduce == "mean":
            mat = arr.mean(axis=0)
        elif args.reduce == "max":
            mat = arr.max(axis=0)
        else:
            sel = max(0, min(int(args.select_idx), arr.shape[0] - 1))
            mat = arr[sel]
    elif arr.ndim == 4:
        # 先在通道维（axis=1）压到单通道，再处理批维
        arr_c = arr.mean(axis=1)  # (B,H,W)
        if args.reduce == "mean":
            mat = arr_c.mean(axis=0)
        elif args.reduce == "max":
            mat = arr_c.max(axis=0)
        else:
            sel = max(0, min(int(args.select_idx), arr_c.shape[0] - 1))
            mat = arr_c[sel]
    else:
        raise ValueError(f"Unsupported array ndim {arr.ndim} for {args.npy}")

    if mat.ndim != 2 or mat.shape[0] != mat.shape[1]:
        raise ValueError(f"Matrix must be square 2D after selection. Got shape {mat.shape} in {args.npy}")
    n = int(mat.shape[0])

    # bins & pixels
    bins_df = build_bins_df(args.chromosome, n, int(args.resolution), int(args.chrom_size))
    pixels_df = build_pixels_df(mat, float(args.min_count_threshold))

    # 写 cooler
    import cooler
    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
    metadata = {
        "chromosome": str(args.chromosome),
        "resolution": int(args.resolution),
        "num_bins": int(n),
        "source": os.path.abspath(args.npy),
        "note": "Converted from N×N .npy matrix; counts stored as float."
    }
    cooler.create_cooler(
        cool_uri=args.output,
        bins=bins_df,
        pixels=pixels_df,
        dtypes={"count": "float64"},
        metadata=metadata,
        symmetric_upper=True,
        ordered=True,
        ensure_sorted=False,
    )
    print(f"Saved cooler -> {args.output}")


if __name__ == "__main__":
    main()


